{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Mathematical Engineering of Deep Learning\n",
    "\n",
    "## Practical 5 (Julia version)\n",
    "**For an R or Python version see the [course website](https://deeplearningmath.org/)**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this practical we deal with a reduced version of the CIFAR10 dataset and train convolutional neural networks. We use a small subset of the dataset to be able to compute things in sensbile time within the tutorial.\n",
    "\n",
    "The main focus is to see how \"to assemble\" different convolutional layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CIFAR10\n",
    "See [CIFAR10 website](https://www.cs.toronto.edu/~kriz/cifar.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "size(CIFAR10Fulltrain_x) = (32, 32, 3, 50000)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10-element Array{String,1}:\n",
       " \"airplane\"\n",
       " \"automobile\"\n",
       " \"bird\"\n",
       " \"cat\"\n",
       " \"deer\"\n",
       " \"dog\"\n",
       " \"frog\"\n",
       " \"horse\"\n",
       " \"ship\"\n",
       " \"truck\""
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using MLDatasets\n",
    "CIFAR10Fulltrain_x, CIFAR10Fulltrain_y = CIFAR10.traindata()\n",
    "CIFAR10Fulltest_x,  CIFAR10Fulltest_y  = CIFAR10.testdata()\n",
    "@show size(CIFAR10Fulltrain_x)\n",
    "classNames = CIFAR10.classnames()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We consider only 3 classes - not all 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length(CIFAR10Filteredtrain_y) = 15000\n",
      "length(CIFAR10Filteredtest_y) = 3000\n"
     ]
    }
   ],
   "source": [
    "#use only \"airplane\", \"cat\" and \"truck\" to make the problem \"easier and quicker\" for the practical\n",
    "usedLabels = [0,3,9]\n",
    "trainFilter = [y ∈ usedLabels for y in CIFAR10Fulltrain_y]  #use \\in +[TAB] to make ∈ (element of set checking)\n",
    "testFilter = [y ∈ usedLabels for y in CIFAR10Fulltest_y]\n",
    "CIFAR10Filteredtrain_x, CIFAR10Filteredtrain_y = CIFAR10Fulltrain_x[:,:,:,trainFilter], CIFAR10Fulltrain_y[trainFilter]\n",
    "CIFAR10Filteredtest_x,  CIFAR10Filteredtest_y  = CIFAR10Fulltest_x[:,:,:,testFilter],  CIFAR10Fulltest_y[testFilter];\n",
    "@show length(CIFAR10Filteredtrain_y)\n",
    "@show length(CIFAR10Filteredtest_y);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We stay with a small training data size and validation data size - also due to time constraints. For testing we could use the full test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainLength = 2000\n",
    "validateLength = 1000;\n",
    "trainRange = 1:trainLength\n",
    "validateRange = (trainLength+1):(trainLength+validateLength);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Creates the training, validation, and test sets \n",
    "using Flux: onehotbatch\n",
    "\n",
    "x_train = CIFAR10Filteredtrain_x[:,:,:,trainRange]\n",
    "x_validate = CIFAR10Filteredtrain_x[:,:,:,validateRange]\n",
    "x_test = CIFAR10Filteredtest_x\n",
    "\n",
    "y_train = onehotbatch(CIFAR10Filteredtrain_y[trainRange], usedLabels)\n",
    "y_validate = onehotbatch(CIFAR10Filteredtrain_y[validateRange], usedLabels)\n",
    "y_test = onehotbatch(CIFAR10Filteredtest_y, usedLabels);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "40"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Setup training mini-batches\n",
    "using Flux: onehotbatch\n",
    "batchSize = 50\n",
    "\n",
    "x_train_batches = [x_train[:, :, :, r] for r in Iterators.partition(trainRange, batchSize)];\n",
    "y_train_batches = [y_train[:,r] for r in Iterators.partition(trainRange,batchSize)];\n",
    "\n",
    "numMiniBatches = length(x_train_batches)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A convolutional model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin with a convolutional model that has two convolutional layers, some batch normalization, maxpooling, and dropout for the dense layers that follow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Chain(Conv((3, 3), 3=>8, relu), BatchNorm(8), #13, Conv((3, 3), 8=>4, relu), BatchNorm(4), #14, flatten, Dense(196, 80, relu), Dropout(0.5), Dense(80, 40, relu), Dropout(0.5), Dense(40, 3), softmax)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Flux\n",
    "\n",
    "function buildModel1()\n",
    "    Chain(  #Assuming 32x32x3 input layer (like CIFAR10)\n",
    "      Conv((3, 3), 3 => 8, relu, pad=(1, 1), stride=(1, 1)),   #32x32x8 convolutional layer\n",
    "      BatchNorm(8),\n",
    "      x -> maxpool(x, (2, 2)),  #16x16x8\n",
    "      Conv((3, 3), 8 => 4, relu, pad=(0, 0), stride=(1, 1)), #14x14x4  convolutional layer\n",
    "      BatchNorm(4),\n",
    "      x -> maxpool(x, (2, 2)), #7x7x4 ,\n",
    "      flatten,  #196 neurons\n",
    "      Dense(196, 80, relu),\n",
    "      Dropout(0.5),\n",
    "      Dense(80, 40, relu), #40 neurons\n",
    "      Dropout(0.5),\n",
    "      Dense(40, 3), #3 output neruons\n",
    "      softmax)\n",
    "end\n",
    "\n",
    "model1 = buildModel1()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To \"debug\" the model, it is good to consider one sample image from the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32×32×3×1 Array{N0f8,4} with eltype FixedPointNumbers.Normed{UInt8,8}:\n",
       "[:, :, 1, 1] =\n",
       " 0.604  0.549  0.549  0.533  0.506  …  0.682  0.643  0.686  0.647  0.639\n",
       " 0.494  0.569  0.545  0.537  0.553     0.553  0.553  0.612  0.612  0.62\n",
       " 0.412  0.49   0.451  0.478  0.533     0.173  0.451  0.604  0.624  0.639\n",
       " 0.4    0.486  0.576  0.518  0.729     0.169  0.443  0.576  0.514  0.569\n",
       " 0.49   0.588  0.541  0.592  0.843     0.224  0.455  0.608  0.369  0.169\n",
       " 0.608  0.596  0.518  0.71   0.792  …  0.204  0.447  0.631  0.4    0.075\n",
       " 0.675  0.682  0.667  0.796  0.643     0.173  0.443  0.627  0.424  0.078\n",
       " 0.706  0.698  0.698  0.816  0.588     0.188  0.455  0.655  0.502  0.29\n",
       " 0.557  0.525  0.671  0.816  0.541     0.31   0.486  0.647  0.604  0.525\n",
       " 0.435  0.431  0.753  0.796  0.467     0.678  0.584  0.596  0.612  0.467\n",
       " 0.416  0.522  0.859  0.702  0.369  …  0.851  0.627  0.639  0.714  0.431\n",
       " 0.427  0.639  0.918  0.663  0.424     0.651  0.557  0.643  0.702  0.388\n",
       " 0.482  0.753  0.898  0.643  0.424     0.38   0.384  0.522  0.49   0.239\n",
       " ⋮                                  ⋱                       ⋮      \n",
       " 0.455  0.557  0.694  0.718  0.522  …  0.31   0.271  0.243  0.137  0.024\n",
       " 0.4    0.376  0.396  0.475  0.49      0.286  0.278  0.2    0.082  0.039\n",
       " 0.373  0.388  0.396  0.357  0.4       0.282  0.271  0.173  0.055  0.098\n",
       " 0.353  0.373  0.345  0.369  0.384     0.329  0.294  0.153  0.043  0.2\n",
       " 0.282  0.349  0.404  0.357  0.345     0.341  0.31   0.169  0.055  0.267\n",
       " 0.235  0.314  0.369  0.302  0.314  …  0.369  0.463  0.4    0.231  0.353\n",
       " 0.22   0.255  0.255  0.271  0.353     0.282  0.337  0.216  0.192  0.455\n",
       " 0.302  0.329  0.325  0.38   0.369     0.169  0.333  0.125  0.212  0.525\n",
       " 0.369  0.361  0.353  0.345  0.271     0.133  0.122  0.09   0.318  0.549\n",
       " 0.357  0.376  0.31   0.298  0.278     0.404  0.302  0.165  0.404  0.561\n",
       " 0.341  0.302  0.267  0.251  0.267  …  0.373  0.267  0.239  0.482  0.561\n",
       " 0.31   0.278  0.263  0.278  0.341     0.486  0.369  0.365  0.514  0.561\n",
       "\n",
       "[:, :, 2, 1] =\n",
       " 0.694  0.627  0.608  0.576  0.537  …  0.643  0.608  0.655  0.604  0.58\n",
       " 0.537  0.6    0.573  0.557  0.573     0.537  0.545  0.604  0.596  0.58\n",
       " 0.408  0.49   0.451  0.475  0.545     0.18   0.467  0.627  0.631  0.612\n",
       " 0.396  0.506  0.6    0.522  0.729     0.161  0.443  0.6    0.51   0.529\n",
       " 0.514  0.631  0.588  0.616  0.863     0.2    0.424  0.6    0.345  0.125\n",
       " 0.651  0.643  0.569  0.757  0.847  …  0.18   0.412  0.604  0.361  0.035\n",
       " 0.745  0.737  0.722  0.871  0.737     0.145  0.424  0.639  0.42   0.055\n",
       " 0.78   0.741  0.741  0.89   0.698     0.145  0.439  0.678  0.506  0.267\n",
       " 0.612  0.545  0.69   0.875  0.635     0.271  0.482  0.647  0.584  0.49\n",
       " 0.471  0.435  0.765  0.859  0.573     0.675  0.612  0.596  0.592  0.431\n",
       " 0.42   0.498  0.855  0.761  0.478  …  0.792  0.612  0.635  0.694  0.396\n",
       " 0.408  0.612  0.914  0.722  0.537     0.537  0.502  0.639  0.686  0.365\n",
       " 0.475  0.753  0.929  0.729  0.545     0.361  0.392  0.541  0.506  0.247\n",
       " ⋮                                  ⋱                       ⋮      \n",
       " 0.459  0.561  0.694  0.718  0.525  …  0.314  0.278  0.251  0.145  0.024\n",
       " 0.396  0.38   0.4    0.482  0.502     0.29   0.29   0.208  0.086  0.035\n",
       " 0.373  0.396  0.404  0.369  0.42      0.275  0.278  0.169  0.047  0.086\n",
       " 0.349  0.376  0.349  0.38   0.408     0.306  0.298  0.141  0.024  0.176\n",
       " 0.275  0.349  0.404  0.369  0.361     0.333  0.322  0.173  0.051  0.251\n",
       " 0.235  0.318  0.373  0.314  0.322  …  0.365  0.475  0.424  0.251  0.353\n",
       " 0.224  0.263  0.263  0.282  0.365     0.251  0.318  0.22   0.192  0.443\n",
       " 0.306  0.337  0.337  0.392  0.38      0.133  0.306  0.102  0.188  0.498\n",
       " 0.376  0.373  0.365  0.357  0.286     0.129  0.106  0.055  0.282  0.51\n",
       " 0.373  0.388  0.322  0.306  0.286     0.424  0.302  0.133  0.365  0.522\n",
       " 0.353  0.314  0.275  0.259  0.275  …  0.384  0.259  0.208  0.447  0.525\n",
       " 0.318  0.286  0.271  0.286  0.349     0.482  0.345  0.325  0.475  0.522\n",
       "\n",
       "[:, :, 3, 1] =\n",
       " 0.733  0.663  0.643  0.608  0.565  …  0.667  0.627  0.651  0.502  0.471\n",
       " 0.533  0.604  0.584  0.573  0.596     0.588  0.596  0.627  0.51   0.478\n",
       " 0.373  0.463  0.439  0.475  0.557     0.22   0.518  0.667  0.557  0.522\n",
       " 0.388  0.518  0.624  0.545  0.745     0.188  0.475  0.639  0.498  0.49\n",
       " 0.545  0.678  0.635  0.655  0.89      0.224  0.459  0.647  0.373  0.125\n",
       " 0.706  0.686  0.604  0.776  0.875  …  0.208  0.451  0.671  0.408  0.047\n",
       " 0.824  0.784  0.745  0.878  0.765     0.173  0.463  0.706  0.471  0.075\n",
       " 0.839  0.769  0.753  0.902  0.725     0.18   0.486  0.729  0.537  0.275\n",
       " 0.612  0.537  0.686  0.882  0.663     0.294  0.525  0.682  0.596  0.478\n",
       " 0.431  0.4    0.741  0.851  0.588     0.659  0.624  0.627  0.604  0.42\n",
       " 0.384  0.471  0.851  0.776  0.522  …  0.761  0.612  0.667  0.702  0.388\n",
       " 0.4    0.612  0.933  0.769  0.604     0.537  0.522  0.675  0.702  0.357\n",
       " 0.459  0.733  0.922  0.745  0.576     0.424  0.455  0.596  0.529  0.243\n",
       " ⋮                                  ⋱                       ⋮      \n",
       " 0.404  0.533  0.698  0.753  0.573  …  0.392  0.361  0.302  0.173  0.043\n",
       " 0.325  0.333  0.38   0.486  0.514     0.361  0.376  0.271  0.118  0.047\n",
       " 0.298  0.329  0.361  0.341  0.396     0.349  0.361  0.224  0.071  0.086\n",
       " 0.31   0.341  0.322  0.361  0.392     0.384  0.376  0.184  0.035  0.165\n",
       " 0.271  0.337  0.388  0.345  0.329     0.412  0.396  0.224  0.078  0.263\n",
       " 0.239  0.302  0.337  0.259  0.259  …  0.416  0.529  0.478  0.302  0.396\n",
       " 0.212  0.235  0.212  0.216  0.286     0.263  0.341  0.251  0.227  0.478\n",
       " 0.282  0.298  0.275  0.314  0.29      0.129  0.306  0.114  0.204  0.522\n",
       " 0.329  0.314  0.294  0.275  0.196     0.145  0.114  0.055  0.286  0.533\n",
       " 0.278  0.306  0.251  0.251  0.235     0.482  0.329  0.141  0.376  0.545\n",
       " 0.278  0.243  0.216  0.208  0.224  …  0.435  0.286  0.224  0.471  0.557\n",
       " 0.275  0.239  0.216  0.231  0.29      0.522  0.369  0.357  0.514  0.565"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sampleImage = x_train_batches[1][:,:,:,1:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×1 Array{Float32,2}:\n",
       " 0.4049066\n",
       " 0.28563324\n",
       " 0.30946016"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#The output of the (untrained) modle on the sample image\n",
    "model1(sampleImage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "length(model1) #The number of `layers`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 1 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(32, 32, 8, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 2 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(32, 32, 8, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 3 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(16, 16, 8, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 4 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(14, 14, 4, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 5 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(14, 14, 4, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 6 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(7, 7, 4, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 7 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(196, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 8 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(80, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 9 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(80, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 10 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(40, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 11 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(40, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 12 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 13 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#This function just debugs layer by layer.\n",
    "function debugModel(model)\n",
    "    for topLayer = 1:length(model)\n",
    "        print(\"Output after $topLayer layers\"); flush(stdout)\n",
    "        outputSize = size(model[1:topLayer](sampleImage))\n",
    "        display(outputSize)\n",
    "    end\n",
    "end\n",
    "debugModel(model1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Task 1**: Modify the model to have (5x5) convolutions in the first laye and no max pooling layer after the second convolutional layer. You'll need to update the number of neurons in the dense layers. The use `debugModel()` to print the size evolution of your model. Call this model, `model2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Chain(Conv((5, 5), 3=>8, relu), BatchNorm(8), #17, Conv((3, 3), 8=>4, relu), BatchNorm(4), flatten, Dense(676, 80, relu), Dropout(0.5), Dense(80, 40, relu), Dropout(0.5), Dense(40, 3), softmax)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#SOLUTION:\n",
    "\n",
    "model2 = Chain(  #Assuming 32x32x3 input layer (like CIFAR10)\n",
    "  Conv((5, 5), 3 => 8, relu, pad=(1, 1), stride=(1, 1)),   #30x3-x8 convolutional layer\n",
    "  BatchNorm(8),\n",
    "  x -> maxpool(x, (2, 2)),  #15x15x8\n",
    "  Conv((3, 3), 8 => 4, relu, pad=(0, 0), stride=(1, 1)), #13x13x4  convolutional layer\n",
    "  BatchNorm(4),\n",
    "  flatten,  #676 neurons\n",
    "  Dense(676, 80, relu),\n",
    "  Dropout(0.5),\n",
    "  Dense(80, 40, relu), #40 neurons\n",
    "  Dropout(0.5),\n",
    "  Dense(40, 3), #3 output neruons\n",
    "  softmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 1 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(30, 30, 8, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 2 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(30, 30, 8, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 3 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(15, 15, 8, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 4 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(13, 13, 4, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 5 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(13, 13, 4, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 6 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(676, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 7 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(80, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 8 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(80, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 9 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(40, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 10 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(40, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 11 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 12 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "debugModel(model2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now train the model using basic functions of `Flux.jl`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "loss (generic function with 1 method)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#define the loss. function\n",
    "using Flux\n",
    "using Flux: crossentropy\n",
    "loss(x, y, model) = crossentropy(model(x), y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.1553065f0"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss(x_train_batches[1],y_train_batches[1],model1) #Loss on first batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.1206558f0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss(x_train,y_train,model1) #Loss on all the training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "accuracy (generic function with 1 method)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#define the accuracy function\n",
    "using Flux: onecold\n",
    "using Statistics\n",
    "accuracy(x, y, model) = mean(onecold(model(x)) .== onecold(y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.338"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy(x_validate,y_validate,model1) #should be around 0.33 as model is garbage at this point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Grads(...)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#example of computing the gradient (automatic differentiantion) on the first minibatch\n",
    "gs = gradient(()->loss(x_train_batches[1],y_train_batches[1],model1),params(model1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Example of a single update step of the optimizer\n",
    "using Flux: update!\n",
    "η = 0.002\n",
    "opt = ADAM(η)\n",
    "update!(opt,params(model1),gs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A training loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we use the above for a simple training loop. On each epoch we print:\n",
    "* The epoch number\n",
    "* The wall time\n",
    "* The accuracy computed on the validation set\n",
    "* The loss (computed on the training set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "trainModel (generic function with 1 method)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Dates\n",
    "function trainModel(model; epochs = 20, η = 0.005)\n",
    "    opt = ADAM(η)\n",
    "    for ep in 1:epochs #Loop over epochs\n",
    "        \n",
    "        for bi in 1:numMiniBatches #Loop over minibatches\n",
    "            gs = gradient(()->loss(x_train_batches[bi],y_train_batches[bi],model),params(model))\n",
    "            update!(opt,params(model),gs)\n",
    "        end\n",
    "        \n",
    "        acc = accuracy(x_validate,y_validate,model)\n",
    "        ls = loss(x_train,y_train,model)\n",
    "        time = Dates.format(now(), \"HH:MM:SS\")\n",
    "        @show ep, time, acc, ls\n",
    "    end\n",
    "    \n",
    "    return model\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This now trains a model of type \"model 1\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(ep, time, acc, ls) = (1, \"22:56:37\", 0.645, 0.85031986f0)\n",
      "(ep, time, acc, ls) = (2, \"22:56:39\", 0.727, 0.67430466f0)\n",
      "(ep, time, acc, ls) = (3, \"22:56:42\", 0.759, 0.58003885f0)\n",
      "(ep, time, acc, ls) = (4, \"22:56:45\", 0.764, 0.5340183f0)\n",
      "(ep, time, acc, ls) = (5, \"22:56:47\", 0.788, 0.48535916f0)\n",
      "(ep, time, acc, ls) = (6, \"22:56:50\", 0.776, 0.4810053f0)\n",
      "(ep, time, acc, ls) = (7, \"22:56:52\", 0.768, 0.41606054f0)\n",
      "(ep, time, acc, ls) = (8, \"22:56:55\", 0.745, 0.51387215f0)\n",
      "(ep, time, acc, ls) = (9, \"22:56:58\", 0.795, 0.376185f0)\n",
      "(ep, time, acc, ls) = (10, \"22:57:00\", 0.791, 0.3298902f0)\n",
      "(ep, time, acc, ls) = (11, \"22:57:03\", 0.785, 0.35687378f0)\n",
      "(ep, time, acc, ls) = (12, \"22:57:06\", 0.781, 0.34605777f0)\n",
      "(ep, time, acc, ls) = (13, \"22:57:08\", 0.801, 0.26681858f0)\n",
      "(ep, time, acc, ls) = (14, \"22:57:11\", 0.814, 0.25501287f0)\n",
      "(ep, time, acc, ls) = (15, \"22:57:13\", 0.79, 0.25198582f0)\n",
      "(ep, time, acc, ls) = (16, \"22:57:16\", 0.786, 0.26506403f0)\n",
      "(ep, time, acc, ls) = (17, \"22:57:19\", 0.812, 0.20088269f0)\n",
      "(ep, time, acc, ls) = (18, \"22:57:21\", 0.782, 0.3156767f0)\n",
      "(ep, time, acc, ls) = (19, \"22:57:24\", 0.801, 0.19398215f0)\n",
      "(ep, time, acc, ls) = (20, \"22:57:26\", 0.821, 0.1730365f0)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Chain(Conv((3, 3), 3=>8, relu), BatchNorm(8), #13, Conv((3, 3), 8=>4, relu), BatchNorm(4), #14, flatten, Dense(196, 80, relu), Dropout(0.5), Dense(80, 40, relu), Dropout(0.5), Dense(40, 3), softmax)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainedModel = trainModel(buildModel1())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "testModel (generic function with 1 method)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#This is the testing function\n",
    "testModel(model) = accuracy(x_test,y_test,model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7953333333333333"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "testModel(trainedModel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Task 2**: Attempt to create a different model archiecture - use two additional convolutional layer. Try to maximize the validation accuracy and we'll see in class who gets the best test accuracy (you can only check the test accuracy once)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Chain(Conv((3, 3), 3=>8, relu), BatchNorm(8), #13, Conv((3, 3), 8=>4, relu), BatchNorm(4), #14, flatten, Dense(196, 80, relu), Dropout(0.5), Dense(80, 40, relu), Dropout(0.5), Dense(40, 3), softmax)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Solution (example)\n",
    "\n",
    "function buildModel3()\n",
    "    Chain(  #Assuming 32x32x3 input layer (like CIFAR10)\n",
    "      Conv((5, 5), 3 => 10, relu, pad=(1, 1), stride=(1, 1)),  \n",
    "      BatchNorm(10),\n",
    "      x -> maxpool(x, (2, 2)), \n",
    "      Conv((3, 3), 10 => 6, relu, pad=(1, 1), stride=(1, 1)), \n",
    "      BatchNorm(6),\n",
    "      x -> maxpool(x, (2, 2)), \n",
    "      Conv((3, 3), 6 => 6, relu, pad=(1, 1), stride=(1, 1)),\n",
    "      BatchNorm(6),\n",
    "      Conv((3, 3), 6 => 6, relu, pad=(1, 1), stride=(1, 1)), \n",
    "      BatchNorm(6),        \n",
    "      flatten, \n",
    "      Dense(294, 120, relu),\n",
    "      Dropout(0.5),\n",
    "      Dense(120, 60, relu),\n",
    "      Dropout(0.5),\n",
    "      Dense(60, 3), #3 output neruons\n",
    "      softmax)\n",
    "end\n",
    "\n",
    "model1 = buildModel1()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 1 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(30, 30, 10, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 2 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(30, 30, 10, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 3 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(15, 15, 10, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 4 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(15, 15, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 5 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(15, 15, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 6 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(7, 7, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 7 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(7, 7, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 8 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(7, 7, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 9 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(7, 7, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 10 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(7, 7, 6, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 11 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(294, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 12 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(120, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 13 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(120, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 14 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(60, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 15 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(60, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 16 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output after 17 layers"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(3, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "debugModel(buildModel3())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(ep, time, acc, ls) = (1, \"22:59:03\", 0.633, 0.90386766f0)\n",
      "(ep, time, acc, ls) = (2, \"22:59:10\", 0.705, 0.6717061f0)\n",
      "(ep, time, acc, ls) = (3, \"22:59:17\", 0.724, 0.56504637f0)\n",
      "(ep, time, acc, ls) = (4, \"22:59:23\", 0.707, 0.5984488f0)\n",
      "(ep, time, acc, ls) = (5, \"22:59:30\", 0.744, 0.49371997f0)\n",
      "(ep, time, acc, ls) = (6, \"22:59:36\", 0.761, 0.46326068f0)\n",
      "(ep, time, acc, ls) = (7, \"22:59:43\", 0.761, 0.47581998f0)\n",
      "(ep, time, acc, ls) = (8, \"22:59:50\", 0.782, 0.4326664f0)\n",
      "(ep, time, acc, ls) = (9, \"22:59:56\", 0.771, 0.39080012f0)\n",
      "(ep, time, acc, ls) = (10, \"23:00:03\", 0.785, 0.38403004f0)\n",
      "(ep, time, acc, ls) = (11, \"23:00:10\", 0.792, 0.34761646f0)\n",
      "(ep, time, acc, ls) = (12, \"23:00:16\", 0.796, 0.33840582f0)\n",
      "(ep, time, acc, ls) = (13, \"23:00:23\", 0.799, 0.30736512f0)\n",
      "(ep, time, acc, ls) = (14, \"23:00:29\", 0.805, 0.2966043f0)\n",
      "(ep, time, acc, ls) = (15, \"23:00:36\", 0.795, 0.2967552f0)\n",
      "(ep, time, acc, ls) = (16, \"23:00:43\", 0.803, 0.2861434f0)\n",
      "(ep, time, acc, ls) = (17, \"23:00:49\", 0.806, 0.26822418f0)\n",
      "(ep, time, acc, ls) = (18, \"23:00:56\", 0.803, 0.24128614f0)\n",
      "(ep, time, acc, ls) = (19, \"23:01:02\", 0.796, 0.2738634f0)\n",
      "(ep, time, acc, ls) = (20, \"23:01:09\", 0.793, 0.23773538f0)\n",
      "(ep, time, acc, ls) = (21, \"23:01:16\", 0.805, 0.22410935f0)\n",
      "(ep, time, acc, ls) = (22, \"23:01:22\", 0.817, 0.18094382f0)\n",
      "(ep, time, acc, ls) = (23, \"23:01:29\", 0.803, 0.18685019f0)\n",
      "(ep, time, acc, ls) = (24, \"23:01:35\", 0.789, 0.24345194f0)\n",
      "(ep, time, acc, ls) = (25, \"23:01:42\", 0.81, 0.16629252f0)\n",
      "(ep, time, acc, ls) = (26, \"23:01:49\", 0.798, 0.15265629f0)\n",
      "(ep, time, acc, ls) = (27, \"23:01:56\", 0.812, 0.14324151f0)\n",
      "(ep, time, acc, ls) = (28, \"23:02:03\", 0.817, 0.17618494f0)\n",
      "(ep, time, acc, ls) = (29, \"23:02:10\", 0.806, 0.13112833f0)\n",
      "(ep, time, acc, ls) = (30, \"23:02:17\", 0.813, 0.114730895f0)\n",
      "(ep, time, acc, ls) = (31, \"23:02:23\", 0.784, 0.23228608f0)\n",
      "(ep, time, acc, ls) = (32, \"23:02:30\", 0.803, 0.09332075f0)\n",
      "(ep, time, acc, ls) = (33, \"23:02:37\", 0.797, 0.12686563f0)\n",
      "(ep, time, acc, ls) = (34, \"23:02:44\", 0.791, 0.13599497f0)\n",
      "(ep, time, acc, ls) = (35, \"23:02:51\", 0.809, 0.09453621f0)\n",
      "(ep, time, acc, ls) = (36, \"23:02:58\", 0.804, 0.07733429f0)\n",
      "(ep, time, acc, ls) = (37, \"23:03:04\", 0.803, 0.08141397f0)\n",
      "(ep, time, acc, ls) = (38, \"23:03:11\", 0.757, 0.3663164f0)\n",
      "(ep, time, acc, ls) = (39, \"23:03:18\", 0.682, 0.9954551f0)\n",
      "(ep, time, acc, ls) = (40, \"23:03:25\", 0.78, 0.17909528f0)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Chain(Conv((5, 5), 3=>10, relu), BatchNorm(10), #29, Conv((3, 3), 10=>6, relu), BatchNorm(6), #30, Conv((3, 3), 6=>6, relu), BatchNorm(6), Conv((3, 3), 6=>6, relu), BatchNorm(6), flatten, Dense(294, 120, relu), Dropout(0.5), Dense(120, 60, relu), Dropout(0.5), Dense(60, 3), softmax)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#training\n",
    "trainedModel = trainModel(buildModel3(), η = 0.001, epochs = 40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.774"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#testing\n",
    "testModel(trainedModel)"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.5.3",
   "language": "julia",
   "name": "julia-1.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
